package ink.mint.neural;

import ink.mint.activation.Activateable;
import ink.mint.activation.ActivationFunction;
import ink.mint.activation.ActivationType;
import ink.mint.loss.LossFunction;
import ink.mint.loss.LossType;
import ink.mint.loss.Lossable;
import ink.mint.matrix.Matrix;
import ink.mint.matrix.Matrixable;
import ink.mint.parameter.Parameter;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;

public class NeuralNetwork implements Serializable {

    /**
     * 序列化版本标识符 serialVersionUID
     */
    private static final long serialVersionUID = 1L;
    /**
     * 网络层集合
     */
    private List<Layer> layerList;
    /**
     * 网络层
     */
    private transient Layer layer;
    /**
     * 损失函数接口
     */
    private transient Lossable lossable;
    /**
     * 损失函数类型
     */
    private transient LossType lossType;
    /**
     * 激活函数接口
     */
    private transient Activateable activateable;
    /**
     * 矩阵接口
     */
    private transient Matrixable matrix;
    /**
     * 输入向量
     */
    private transient double[] input;
    /**
     * 期望输出向量
     */
    private transient double[] expectations;
    /**
     * 损失值
     */
    private transient double lossValue;
    /**
     * 损失函数求导值
     */
    private transient double[] lossVactor;
    /**
     * 学习率
     */
    private transient double learningRate;
    /**
     * 数据集Xi转换的矩阵 图片
     */
    private transient double[][] X;
    /**
     * 数据集Yi转换的向量 标签
     */
    private transient double[] Y;

    /**
     * 正确数
     */
    transient double correct = 0;
    /**
     * 正确率
     */
    transient double correctRate = 0;
    /**
     * 错误数
     */
    transient double mistake = 0;

    /**
     * 网络构造函数
     */
    public NeuralNetwork() {
        this.layerList = new ArrayList<>();
        this.lossable = new LossFunction();
        this.activateable = new ActivationFunction();
        this.matrix = new Matrix();
    }

    /**
     * 神经网络设置数据集图片 Xi
     *
     * @param X 数据集Yi转换的向量 标签
     * @return 神经网络
     */
    public NeuralNetwork setX(double[][] X) {
        this.X = X;
        return this;
    }

    /**
     * 神经网络设置数据集标签 Yi
     *
     * @param Y 数据集Yi转换的向量 标签
     * @return 神经网络
     */
    public NeuralNetwork setY(double[] Y) {
        this.Y = Y;
        return this;
    }

    /**
     * 神经网络输入函数
     *
     * @param input 输入向量
     * @return 神经网络
     */
    private NeuralNetwork input(double[] input) {
        this.input = input;
        return this;
    }

    /**
     * 神经网络期望值函数
     *
     * @param expectations 期望输出向量
     * @return 神经网络
     */
    public NeuralNetwork expectations(double[] expectations) {
        this.expectations = expectations;
        return this;
    }

    /**
     * 神经网络添加网络层函数
     *
     * @param layer 网络层
     * @return 神经网络
     */
    public NeuralNetwork addLayer(Layer layer) {
        layerList.add(layer);
        return this;
    }

    /**
     * 神经网络创造网络
     *
     * @param isTrain 是否是训练模式（true: 训练 | false: 运行）
     * @return 神经网络
     */
    public NeuralNetwork createNetwork(boolean isTrain) {
        Layer previousLayer;
        Layer nextLayer;
        for (int i = 0; i < layerList.size(); i++) {
            nextLayer = layerList.get(i);
            if (i == 0) {
                nextLayer.input(input, isTrain);
            } else {
                previousLayer = layerList.get(i - 1);
                double[] prevActivateLayer = previousLayer.getActivateLayer();
                nextLayer.input(prevActivateLayer, isTrain);
            }
            nextLayer.activate(nextLayer.getActivationType());
        }
        return this;
    }

    /**
     * 神经网络设置权重weight和偏置bias
     *
     * @param layerIndex 网络层索引
     * @param newWeight  新的权重weight
     * @param newBias    新的偏置bias
     * @return 神经网络
     */
    public NeuralNetwork setWeightAndBias(int layerIndex, double[][] newWeight, double[] newBias) {
        Layer layer = layerList.get(layerIndex);
        layer.setWeight(newWeight);
        layer.setBias(newBias);
        return this;
    }

    /**
     * 神经网络设置损失函数类型
     *
     * @param lossType 损失函数类型
     * @return 神经网络
     */
    public NeuralNetwork setLoss(LossType lossType) {
        this.lossType = lossType;
        return this;
    }

    /**
     * 神经网络设置学习率
     *
     * @param learningRate 学习率
     * @return 神经网络
     */
    public NeuralNetwork setLearningRate(double learningRate) {
        this.learningRate = learningRate;
        return this;
    }

    /**
     * 神经网络计算损失函数及求导
     *
     * @return 神经网络
     */
    public NeuralNetwork lossFunction() {
        Layer endLayer = layerList.get(layerList.size() - 1);
        double[] outputLayer = endLayer.getActivateLayer();
        lossable.setLayer(endLayer);
        this.lossValue = lossable.LossSolve(expectations, outputLayer, lossType);
        this.lossVactor = lossable.derivationLossSolve(expectations, outputLayer, lossType);
        return this;
    }

    transient double[] zLayer;
    transient double[] xLayer;
    transient ActivationType activationType;
    transient double[] derivationActivationVector;
    transient double[][] deltaE;
    transient double[][] result;
    transient double[][] daltaWeight;

    /**
     * 神经网络权重和偏置误差计算函数
     *
     * @param layerIndex 网络层索引
     */
    private void deltaWeightMatrixAndBiasVector(int layerIndex) {
        layer = layerList.get(layerIndex);
        if (layerIndex == layerList.size() - 1) {
            zLayer = layer.getzLayer();
            xLayer = layer.getInput();
            activationType = layer.getActivationType();
            derivationActivationVector = activateable.derivationActivationSolve(zLayer, activationType);
            deltaE = matrix.elementProduct(matrix.vectorToColumnMatrix(lossVactor), matrix.vectorToColumnMatrix(derivationActivationVector));
            result = matrix.multi(deltaE, matrix.vectorToRowMatrix(xLayer));
            daltaWeight = matrix.elementProduct(matrix.customNumber(learningRate, result.length, result[0].length), result);
            layer.setDeltaWeight(daltaWeight);
            layer.setDeltaBias(matrix.matrixToVector(deltaE));
        } else {
            Layer nextLayer = layerList.get(layerIndex + 1);
            double[][] deltaWeight = nextLayer.getDeltaWeight();
            double[][] weight = nextLayer.getWeight();
            double[][] delta = matrix.multi(matrix.transpose(weight), deltaWeight);
            zLayer = layer.getzLayer();
            xLayer = layer.getInput();
            activationType = layer.getActivationType();
            derivationActivationVector = activateable.derivationActivationSolve(zLayer, activationType);
            deltaE = matrix.multi(delta, matrix.vectorToColumnMatrix(derivationActivationVector));
            result = matrix.multi(deltaE, matrix.vectorToRowMatrix(xLayer));
            daltaWeight = matrix.elementProduct(matrix.customNumber(learningRate, result.length, result[0].length), result);
            layer.setDeltaWeight(daltaWeight);
            layer.setDeltaBias(matrix.matrixToVector(deltaE));
        }
    }

    /**
     * 神经网络更新
     *
     * @param layerIndex 网络层索引
     */
    private void update(int layerIndex) {
        Layer layer = layerList.get(layerIndex);
        double[][] newWeight = matrix.sub(layer.getWeight(), layer.getDeltaWeight());
        double[] newBias = matrix.sub(layer.getBias(), layer.getDeltaBias());
        setWeightAndBias(layerIndex, newWeight, newBias);
    }

    /**
     * 神经网络梯度下降函数
     *
     * @return 神经网络
     */
    public NeuralNetwork gradientDescent() {
        int lr = layerList.size();
        for (int i = lr - 1; i >= 0; i--) {
            deltaWeightMatrixAndBiasVector(i);
            update(i);
        }
        return this;
    }

    /**
     * 神经网络训练函数
     *
     * @return 神经网络
     */
    public NeuralNetwork train() {
        int xr = X.length;
        int dimensionality = layerList.get(layerList.size() - 1).getNeuralNumber();
        for (int i = 0; i < xr; i++) {
            input(X[i]);
            expectations(matrix.oneHotEncoding((int) Y[i], dimensionality));
            createNetwork(true);
            lossFunction();
            output(i);
            gradientDescent();
        }
        return this;
    }

    /**
     * 神经网络运行函数
     *
     * @return 神经网络
     */
    public NeuralNetwork run() {
        int xr = X.length;
        int dimensionality = layerList.get(layerList.size() - 1).getNeuralNumber();
        for (int i = 0; i < xr; i++) {
            input(X[i]);
            expectations(matrix.oneHotEncoding((int) Y[i], dimensionality));
            createNetwork(false);
            output(i);
        }
        return this;
    }

    public double[] output() {
        layer = layerList.get(layerList.size() - 1);
        return layer.getActivateLayer();
    }

    /**
     * 神经网络输出函数
     *
     * @param i 索引
     * @return 神经网络
     */
    public NeuralNetwork output(int i) {
        double[] result;
        int maxIndex;
        layer = layerList.get(layerList.size() - 1);
        result = layer.getActivateLayer();
        maxIndex = matrix.maxVectorIndex(result);

        double YIndex = Y[i];

        if (maxIndex == YIndex) {
            correct++;

        } else {
            mistake++;
        }
        correctRate = correct / (correct + mistake);
        System.out.println("第 " + i + " 个 " + "网络输出:  " + maxIndex + "  " + " 期望：" + YIndex + " | 正确数： " + correct + " | 错误数： " + mistake + " | 损失: " + lossValue + " | 正确率： " + correctRate);
        return this;
    }

    /**
     * 获取网络层集合
     *
     * @return 网络层集合
     */
    public List<Layer> getLayerList() {
        return layerList;
    }

    /**
     * 设置网络层集合
     *
     * @param layerList 网络层集合
     */
    public void setLayerList(List<Layer> layerList) {
        this.layerList = layerList;
    }

    /**
     * 获取网络层
     *
     * @return 网络层
     */
    public Layer getLayer() {
        return layer;
    }

    /**
     * 设置网络层
     *
     * @param layer 网络层
     */
    public void setLayer(Layer layer) {
        this.layer = layer;
    }

    /**
     * 获取损失函数接口
     *
     * @return 损失函数接口
     */
    public Lossable getLossable() {
        return lossable;
    }

    /**
     * 设置损失函数接口
     *
     * @param lossable 损失函数接口
     */
    public void setLossable(Lossable lossable) {
        this.lossable = lossable;
    }

    /**
     * 获取损失函数类型
     *
     * @return 损失函数类型
     */
    public LossType getLossType() {
        return lossType;
    }

    /**
     * 设置损失函数类型
     *
     * @param lossType 损失函数类型
     */
    public void setLossType(LossType lossType) {
        this.lossType = lossType;
    }

    /**
     * 获取激活函数接口
     *
     * @return 激活函数接口
     */
    public Activateable getActivateable() {
        return activateable;
    }

    /**
     * 设置激活函数接口
     *
     * @param activateable 激活函数接口
     */
    public void setActivateable(Activateable activateable) {
        this.activateable = activateable;
    }

    /**
     * 获取矩阵接口
     *
     * @return 矩阵接口
     */
    public Matrixable getMatrix() {
        return matrix;
    }

    /**
     * 设置矩阵接口
     *
     * @param matrix 矩阵接口
     */
    public void setMatrix(Matrixable matrix) {
        this.matrix = matrix;
    }

    /**
     * 获取输入向量
     *
     * @return 输入向量
     */
    public double[] getInput() {
        return input;
    }

    /**
     * 设置输入向量
     *
     * @param input 输入向量
     */
    public void setInput(double[] input) {
        this.input = input;
    }

    /**
     * 获取期望向量
     *
     * @return 期望向量
     */
    public double[] getExpectations() {
        return expectations;
    }

    /**
     * 设置期望向量
     *
     * @param expectations 期望向量
     */
    public void setExpectations(double[] expectations) {
        this.expectations = expectations;
    }

    /**
     * 获取损失函数值
     *
     * @return 损失函数值
     */
    public double getLossValue() {
        return lossValue;
    }

    /**
     * 设置损失函数值
     *
     * @param lossValue 损失函数值
     */
    public void setLossValue(double lossValue) {
        this.lossValue = lossValue;
    }

    /**
     * 获取损失函数求导后的向量
     *
     * @return 损失函数求导后的向量
     */
    public double[] getLossVactor() {
        return lossVactor;
    }

    /**
     * 设置损失函数求导后的向量
     *
     * @param lossVactor 损失函数求导后的向量
     */
    public void setLossVactor(double[] lossVactor) {
        this.lossVactor = lossVactor;
    }

    /**
     * 获取学习率
     *
     * @return 学习率
     */
    public double getLearningRate() {
        return learningRate;
    }

    /**
     * 获取数据集Xi转换的矩阵 图片
     *
     * @return 数据集Xi转换的矩阵 图片
     */
    public double[][] getX() {
        return X;
    }

    /**
     * 数据集Yi转换的向量 标签
     *
     * @return 数据集Yi转换的向量 标签
     */
    public double[] getY() {
        return Y;
    }
}
